from core.cuda import cuda, cuda_module
from core.function import Function

from utils.common import pair

from utils.functions_conv import get_deconv_outsize, col2im_array
import utils
import implement

class Deconv2d(Function):
    def __init__(self, stride=1, pad=0, outsize=None):
        """
        Deconv2d 操作的函数类，用于执行反卷积操作。

        Args:
            stride (int or tuple of int): 步幅。
            pad (int or tuple of int): 填充。
            outsize (tuple of int, optional): 输出尺寸。默认为 None。
        """
        super().__init__()
        self.stride = pair(stride)
        self.pad = pair(pad)
        self.outsize = outsize

    def forward(self, x, W, b):
        """
        前向传播，执行反卷积操作。

        Args:
            x (Variable): 输入变量。
            W (Variable): 卷积核变量。
            b (Variable, optional): 偏置变量。默认为 None。

        Returns:
            Variable: 输出变量。
        """
        xp = cuda_module

        Weight = W
        SH, SW = self.stride
        PH, PW = self.pad
        C, OC, KH, KW = Weight.shape
        N, C, H, W = x.shape
        if self.outsize is None:
            out_h = get_deconv_outsize(H, KH, SH, PH)
            out_w = get_deconv_outsize(W, KW, SW, PW)
        else:
            out_h, out_w = pair(self.outsize)
        img_shape = (N, OC, out_h, out_w)

        gcol = xp.tensordot(Weight, x, (0, 1))
        gcol = xp.rollaxis(gcol, 3)
        y = col2im_array(gcol, img_shape, (KH, KW), self.stride, self.pad,
                         to_matrix=False)
        # b, k, h, w
        if b is not None:
            b = cuda.to_array(b)
            self.no_bias = True
            y += b.reshape((1, b.size, 1, 1))
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量和卷积核的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            tuple: 包含输入变量和卷积核的梯度的元组。
        """

        x, W, b = self.inputs

        # ==== gx ====
        gx = utils.functions_collect.conv2d(gy, W, b=None, stride=self.stride, pad=self.pad)
        # ==== gW ====
        f = implement.functions.conv.conv2d_gradw.Conv2DGradW(self)
        gW = f(gy, x)
        # ==== gb ====
        gb = None
        if b.data is not None:
            gb = gy.sum(axis=(0, 2, 3))
        return gx, gW, gb

    def deconv2d(self, x, W, b=None):
        """
        执行反卷积操作。

        Args:
            x (Variable): 输入变量。
            W (Variable): 卷积核变量。
            b (Variable, optional): 偏置变量。默认为 None。

        Returns:
            Variable: 输出变量。
        """
        return self(x, W, b)
